from typing import Dict, Optional
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import OAuth2PasswordBearer
import httpx
from jose import JWTError, jwt
import random  # 用于简单负载均衡（轮询）
from .config import ALLOWED_METHODS
from common.jwt_utils import verify_token, TokenPayload, create_access_token, pwd_context
from common.consul_utils import consul_service  # 引入Consul工具

app = FastAPI(title="API Gateway with Service Discovery")

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# JWT认证相关（同之前）
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/token")

async def get_current_user(token: str = Depends(oauth2_scheme)) -> TokenPayload:
    credentials_exception = HTTPException(
        status_code=401,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        return verify_token(token)
    except JWTError:
        raise credentials_exception

WHITELIST_PATHS = ["/api/users", "/api/auth/token"]

async def jwt_auth_middleware(
    request: Request,
    current_user: TokenPayload = Depends(get_current_user)
):
    path = request.url.path
    for white_path in WHITELIST_PATHS:
        if path.startswith(white_path):
            return
    return current_user

# 服务路由映射（服务名 -> 客户端请求前缀）
SERVICE_ROUTES = {
    "user-service": "/api/users",    # 客户端访问/api/users -> 转发到user-service
    "friend-service": "/api/friends" # 客户端访问/api/friends -> 转发到friend-service
}

def get_service_instance(service_name: str) -> Optional[Dict]:
    """从Consul获取服务实例（简单轮询负载均衡）"""
    instances = consul_service.discover_service(service_name)
    if not instances:
        return None
    # 简单轮询：随机选择一个实例（实际可按权重/负载选择）
    return random.choice(instances)

async def get_http_client():
    async with httpx.AsyncClient() as client:
        yield client

@app.api_route("/{path:path}", methods=ALLOWED_METHODS)
async def forward_request(
    request: Request,
    path: str,
    client: httpx.AsyncClient = Depends(get_http_client),
    auth_result: Optional[TokenPayload] = Depends(jwt_auth_middleware)
):
    request_path = f"/{path}"
    
    # 1. 匹配服务名（根据客户端请求路径前缀）
    target_service = None
    service_prefix = None
    for service, prefix in SERVICE_ROUTES.items():
        if request_path.startswith(prefix):
            target_service = service
            service_prefix = prefix
            break
    
    if not target_service:
        raise HTTPException(status_code=404, detail="Route not found in gateway")
    
    # 2. 从Consul获取服务实例
    instance = get_service_instance(target_service)
    if not instance:
        raise HTTPException(
            status_code=503,
            detail=f"Service {target_service} is unavailable"
        )
    
    # 3. 重写路径并转发请求
    target_host = instance["host"]
    target_port = instance["port"]
    rewritten_path = request_path.replace(service_prefix, "", 1)  # 移除客户端前缀
    target_url = f"http://{target_host}:{target_port}{rewritten_path}"

    # 4. 处理请求头
    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("origin", None)
    if auth_result:
        headers["X-User-ID"] = str(auth_result.user_id)

    # 5. 转发请求
    try:
        response = await client.request(
            method=request.method,
            url=target_url,
            headers=headers,
            content=await request.body(),
            params=dict(request.query_params)
        )
    except httpx.HTTPError as e:
        raise HTTPException(
            status_code=503,
            detail=f"Failed to connect to {target_service}: {str(e)}"
        )

    return Response(
        content=response.content,
        status_code=response.status_code,
        headers=dict(response.headers),
        media_type=response.media_type
    )

# 登录接口（从Consul获取user-service地址）
@app.post("/api/auth/token")
async def login_for_access_token(
    username: str,
    password: str,
    client: httpx.AsyncClient = Depends(get_http_client)
):
    # 从Consul获取user-service实例
    user_instance = get_service_instance("user-service")
    if not user_instance:
        raise HTTPException(status_code=503, detail="User service unavailable")
    
    # 调用user-service查询用户
    try:
        user_url = f"http://{user_instance['host']}:{user_instance['port']}/users/by-username/{username}"
        response = await client.get(user_url)
        user = response.json()
    except httpx.HTTPError:
        raise HTTPException(status_code=500, detail="Failed to connect to user service")

    # 验证密码并返回令牌
    if not user or not pwd_context.verify(password, user["hashed_password"]):
        raise HTTPException(status_code=401, detail="Incorrect username or password")
    
    return {
        "access_token": create_access_token(data={"sub": user["id"]}),
        "token_type": "bearer"
    }
